import pyrtl
import or1200_definitions

width = or1200_definitions.OR1200_OPERAND_WIDTH

class Or1200Mem2Reg(object):
    def __init__(self):
        self.addr = pyrtl.WireVector(bitwidth=2, name='addr')
        self.lsu_op = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_LSUOP_WIDTH)
        self.memdata = pyrtl.WireVector(bitwidth=width, name='memdata')
        # self.addr = pyrtl.Input(bitwidth=2, name='addr')
        # self.lsu_op = pyrtl.Input(bitwidth=or1200_definitions.OR1200_LSUOP_WIDTH)
        # self.memdata = pyrtl.Input(bitwidth=width, name='memdata')
        # self.regdata = pyrtl.Output(bitwidth=width, name='regdata')

        # OR1200_IMPL_MEM2REG2 is not defined
        self.regdata = pyrtl.WireVector(bitwidth=width, name='regdata')
        self.aligned = pyrtl.WireVector(bitwidth=width, name='aligned')

        # initialize all the parts
        alignedClass = AlignedClass()
        bytesClass = BytesClass()

        # establish connection relations for Aligned
        self.aligned <<= alignedClass.aligned
        alignedClass.addr <<= self.addr
        alignedClass.memdata <<= self.memdata

        # establish connection relations for Bytes
        self.regdata <<= bytesClass.regdata
        bytesClass.lsu_op <<= self.lsu_op
        bytesClass.aligned.next <<= self.aligned



# Aligned
class AlignedClass(object):
    def __init__(self):
        self.addr = pyrtl.WireVector(bitwidth=2)
        self.aligned = pyrtl.Register(bitwidth=width)
        self.memdata = pyrtl.WireVector(bitwidth=width)
        with pyrtl.conditional_assignment:
            with self.addr.__eq__(pyrtl.Const(0b00, bitwidth=2)):
                self.aligned.next |= self.memdata
            with self.addr.__eq__(pyrtl.Const(0b00, bitwidth=2)):
                self.aligned.next |= pyrtl.concat(self.memdata[0:24], pyrtl.Const(0b0, bitwidth=8))
            with self.addr.__eq__(pyrtl.Const(0b00, bitwidth=2)):
                self.aligned.next |= pyrtl.concat(self.memdata[0:16], pyrtl.Const(0b0, bitwidth=16))
            with self.addr.__eq__(pyrtl.Const(0b00, bitwidth=2)):
                self.aligned.next |= pyrtl.concat(self.memdata[0:7], pyrtl.Const(0b0, bitwidth=24))

# Bytes
class BytesClass(object):
    def __init__(self):
        self.lsu_op = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_LSUOP_WIDTH)
        self.regdata = pyrtl.WireVector(bitwidth=width)
        self.aligned = pyrtl.Register(bitwidth=width)

        with pyrtl.conditional_assignment:
            with self.lsu_op.__eq__(or1200_definitions.OR1200_LSUOP_LBZ):
                self.regdata |= pyrtl.concat(self.aligned[24:], pyrtl.Const(0b0, bitwidth=24))
            with self.lsu_op.__eq__(or1200_definitions.OR1200_LSUOP_LBS):
                self.regdata |= pyrtl.concat(self.aligned[24:], self.aligned[31]*24)
            with self.lsu_op.__eq__(or1200_definitions.OR1200_LSUOP_LHZ):
                self.regdata |= pyrtl.concat(self.aligned[16:], pyrtl.Const(0b0, bitwidth=16))
            with self.lsu_op.__eq__(or1200_definitions.OR1200_LSUOP_LHS):
                self.regdata |= pyrtl.concat(self.aligned[16:], self.aligned[31]*16)
            with pyrtl.otherwise:
                self.regdata |= self.aligned

if __name__ == '__main__':
    mem2reg = Or1200Mem2Reg()
    ff = open("D:\work_doc\PyRTL\examples\Opencore\or1200_mem2reg.c", 'w')
    essent = pyrtl.ESSENT(CP=40)
    essent._create_code(lambda s: ff.write(s + '\n'))
    ff.close()
    ff = open("D:\work_doc\PyRTL\examples\Opencore\or1200_mem2reg_com.c", 'w')
    com = pyrtl.CompiledSimulation()
    com._create_code(lambda s: ff.write(s + '\n'))
    ff.close()

    best_random_random = 0
    best_select_random = 0
    best_select_select = 0

    input_width = []
    from pyrtl.wire import Input

    IN = pyrtl.working_block().wirevector_subset(Input)
    for inn in IN:
        name = inn.name
        width = inn.bitwidth
        input_width.append([name, width])
    import random

    num = 500
    for i in range(num):
        print(i)
        inputs_random = {}
        inputs_select = {}
        inputs_old = {}
        inputs_ = {}

        for inp in input_width:
            inputs_random[inp[0]] = []
            inputs_select[inp[0]] = []
            result = random.randint(0, 2 ** inp[1] - 1)
            inputs_old[inp[0]] = result

        for j in range(10000):
            for inp in input_width:
                result = random.randint(0, 2 ** inp[1] - 1)
                inputs_[inp[0]] = result
                inputs_random[inp[0]].append(result)
            if random.random() < 0.2:
                for name in inputs_old:
                    inputs_select[name].append(inputs_old[name])
            else:
                for name in inputs_:
                    inputs_select[name].append(inputs_[name])
                    inputs_old[name] = inputs_[name]

        time1 = essent.step_multiple(inputs_random)
        time2 = essent.step_multiple(inputs_select)
        time3 = com.step_multiple(inputs_random)
        time4 = com.step_multiple(inputs_select)

        best_random_random += time3 / time1
        best_select_random += time1 / time2
        best_select_select += time4 / time2
    print(best_random_random / num)
    print(best_select_random / num)
    print(best_select_select / num)
